#include "core/memory.h"
#include "cpu/mmu.h"
#include "dev/console.h"
#include "tools/log.h"
#include "tools/klib.h"

static addr_alloc_t paddr_alloc;        // 物理地址分配结构，内核的位图存放分配器
// 4096的整数倍对齐
// 内核页目录表
static pde_t kernel_page_dir[PDE_CNT]   __attribute__((aligned(MEM_PAGE_SIZE))) ;           // 系统的一级页表，写入到cr3寄存器里

static void addr_alloc_init (addr_alloc_t * alloc, uint8_t * bits,
                    uint32_t start, uint32_t size, uint32_t page_size) {
    //初始化内存分配的结构, 一个page_size 对应的addr_alloc_t里的位图的一位
    alloc->start=start;
    alloc->size=size;
    alloc->page_size=page_size;
    
    mutex_init(&alloc->mutex);
    bitmap_init(&alloc->bitmap,bits,size/page_size,0);
}

/**
 * @brief 分配多页内存
 */
static uint32_t addr_alloc_page (addr_alloc_t * alloc, int page_count) {
    //在位图里找page_count个空闲的位置，返回空闲位置的起始地址
    mutex_lock(&alloc->mutex);
    uint32_t addr=0;
    uint32_t page_index= bitmap_alloc_nbits(&alloc->bitmap,0,page_count);
    if(page_index >=0 ){
        addr=alloc->start + (page_index * alloc->page_size);
    }

    mutex_unlock(&alloc->mutex);

    // 返回的是物理地址，还需要做一个和虚拟地址的映射
    return addr;
}

/**
 * @brief 释放多页内存
 */
static void addr_free_page (addr_alloc_t * alloc, uint32_t addr, int page_count) {
    //释放addr之后的page_count个页，位图置0即可
    mutex_lock(&alloc->mutex);

    uint32_t page_index=(addr-alloc->start)/(alloc->page_size);
    bitmap_set_bit(&alloc->bitmap,page_index,page_count,0);

    mutex_unlock(&alloc->mutex);
}

static void show_mem_info (boot_info_t * boot_info) {
    // 打印内存信息，启动的时候，硬件传过来的信息，可用内存是多少
    log_printf("mem region:");
    for (int i = 0; i < boot_info->ram_region_count; i++) {
        log_printf("[%d]: 0x%x - 0x%x", i,
                    boot_info->ram_region_cfg[i].start,
                    boot_info->ram_region_cfg[i].size);
    }
    log_printf("\n");
}

/**
 * @brief 获取可用的物理内存大小
 */
static uint32_t total_mem_size(boot_info_t * boot_info) {
    int mem_size = 0;

    // 简单起见，暂不考虑中间有空洞的情况
    for (int i = 0; i < boot_info->ram_region_count; i++) {
        mem_size += boot_info->ram_region_cfg[i].size;
    }
    return mem_size;
}

pte_t * find_pte(pde_t * page_dir,uint32_t vaddr , int alloc){
    // 在pde里找vaddr 
    // alloc 为1 未找到的时候分配   alloc为0  未找到的时候不分配
    pde_t *pde=page_dir + pde_index(vaddr);

    if(pde->present){
        // 页表中存在这一项
        return (pte_t *)(pde_paddr(pde)) + pte_index(vaddr);
    }else{
        //不需要分配
        if(alloc==0){
            return (pte_t *)0;
        }else{
            // 分配一个物理页表
            uint32_t pg_paddr = addr_alloc_page(&paddr_alloc, 1);
            if (pg_paddr == 0) {
                return (pte_t *)0;
            }

            // 设置为用户可读写，将被pte中设置所覆盖
            // pg_paddr已经是4096对齐, 后12位为0
            pde->v = pg_paddr | PTE_P | PTE_W | PDE_U;

            // 初始化二级页表
            kernel_memset((void * )pg_paddr,0,MEM_PAGE_SIZE);
            return (pte_t *)(pg_paddr) + pte_index(vaddr);

        }


    }

}

int memory_create_map(pde_t *page_dir , uint32_t vaddr , uint32_t paddr, int count, uint32_t perm){
    // 为一级页表page_dir添加一个映射，将page_dir[i]指向新的页表
    for(int i=0;i<count;i++){
        pte_t *pte=find_pte(page_dir,vaddr,1);
        if(pte==(pte_t *)0){
            return -1;
        }
        
       
        ASSERT(pte->present==0);
        pte->v=paddr | perm | PTE_P; // 写入物理地址、特权级
        // 验证虚拟内存和物理内存的映射关系
        // pte_t * my_pte=(pte_t * )pde_paddr(page_dir + pde_index(vaddr));
        // my_pte=my_pte+pte_index(vaddr);
        // uint32_t cal_addr=pte_paddr(my_pte)+(vaddr & 0xfff);
        // log_printf("virutal 0x%x , physical 0x%x ,calcu 0x%x",vaddr,paddr,cal_addr);
        // 向下, 写入下一个块

        // log_printf("virutal 0x%x , physical 0x%x maps",vaddr,paddr);
        vaddr+=MEM_PAGE_SIZE;
        paddr+=MEM_PAGE_SIZE;

    }
};

static void create_kernel_table(){
    //创建内核页表
    extern uint8_t s_text[], e_text[], s_data[], e_data[];
    extern uint8_t kernel_base[];

    // 地址映射表, 用于建立内核级的地址映射
    // 地址不变，但是添加了属性
    static memory_map_t kernel_map[] = {
        {kernel_base,   s_text,         0,              PTE_W},         // 内核栈区
        {s_text,        e_text,         s_text,         0},         // 内核代码区
        {s_data,        (void *)(MEM_EBDA_START),   s_data,        PTE_W},      // 内核数据区
        {(void *)CONSOLE_DISP_ADDR, (void *)(CONSOLE_DISP_END - 1), (void *)CONSOLE_VIDEO_BASE, PTE_W}, //显存的区域
        {(void *)(MEM_EXT_START) , (void *)(MEM_EXT_END), (void *)(MEM_EXT_START) , PTE_W}  // 操作系统为进程分配页表的位置
        // 扩展存储空间一一映射，方便直接操作
       // {(void *)MEM_EXT_START, (void *)MEM_EXT_END,     (void *)MEM_EXT_START, PTE_W},
    };
    //内核各个数据段的值，由gcc链接脚本获得，获取各个段的起始地址
    //虚拟地址和实际地址是一样的。程序在打开分页之前，eip是直接指向物理地址的，打开之后，需要内核的虚拟地址和物理地址一样
    
    // 清空页目录表
    kernel_memset(kernel_page_dir, 0, sizeof(kernel_page_dir));


    //遍历内核的每一个段，为其建立映射关系
    for(int i=0;i<sizeof(kernel_map) / sizeof(memory_map_t);i++){
        memory_map_t * map = kernel_map + i;

        // 可能有多个页，建立多个页的配置
        // 简化起见，不考虑4M的情况
        int vstart = down2((uint32_t)map->vstart, MEM_PAGE_SIZE);
        int vend = up2((uint32_t)map->vend, MEM_PAGE_SIZE);
        int page_count = (vend - vstart) / MEM_PAGE_SIZE;

        memory_create_map(kernel_page_dir, vstart, (uint32_t)map->pstart, page_count, map->perm);
   
    }
}
void memory_init (boot_info_t * boot_info) {
    //内存初始化

    // 1MB内存空间起始，在链接脚本中定义
    // 放在操作系统.bss节之后，因为不确定运行时内存的大小，就不能用全局变量来定义
    extern uint8_t * mem_free_start;
    log_printf("mem init.\n");
    show_mem_info(boot_info);

    // 在内核数据后面放物理页位图
    uint8_t * mem_free = (uint8_t *)&mem_free_start;  
    

    // 计算1MB以上空间的空闲内存容量，并对齐的页边界
    uint32_t mem_up1MB_free = total_mem_size(boot_info) - MEM_EXT_START;
    mem_up1MB_free = down2(mem_up1MB_free, MEM_PAGE_SIZE);   // 对齐到4KB页
    log_printf("Free memory: 0x%x, size: 0x%x\n", MEM_EXT_START, mem_up1MB_free);

    // 4GB大小需要总共4*1024*1024*1024/4096/8=128KB的位图, 使用低1MB的RAM空间中足够
    // 该部分的内存仅跟在mem_free_start开始放置
    addr_alloc_init(&paddr_alloc, mem_free, MEM_EXT_START, mem_up1MB_free, MEM_PAGE_SIZE);
    mem_free += bitmap_byte_count(paddr_alloc.size / MEM_PAGE_SIZE);

    // 到这里，mem_free应该比EBDA地址要小
    ASSERT(mem_free < (uint8_t *)MEM_EBDA_START);

    create_kernel_table();
    //将内核页表写入cr3
    mmu_set_page_dir((uint32_t)kernel_page_dir);
}

uint32_t memory_create_uvm (void){
    //为进程创建一级页表,返回新的一级页表地址，出错返回0
    pde_t * pde_paddr =(pde_t *)addr_alloc_page(&paddr_alloc,1);
    if(!pde_paddr){
        //创建页表失败
        return 0;
    }

    //清空数据
    kernel_memset((void *)pde_paddr,0,MEM_PAGE_SIZE);

    //将操作系统的数据搬过来
    // 这个是取操作系统最大内存的页表项
    uint32_t user_pde_start=pde_index(MEMORY_TASK_BASE);
    for(int i=0;i<user_pde_start;i++){
        pde_paddr[i].v=kernel_page_dir[i].v;
    }
    return (uint32_t )pde_paddr;

};

uint32_t memory_alloc_page_for (uint32_t addr, uint32_t size, int perm){
    return memory_alloc_for_page_dir(get_task_current()->tss.cr3,addr,size,perm);
};

uint32_t memory_alloc_for_page_dir (uint32_t page_dir, uint32_t vaddr, uint32_t size, int perm){
    //为进程分配页表，page_dir是一级页表的地址，vaddr是虚拟地址，size是大小，perm是权限
    //log_printf("try to alloc vaddr to paddr page_dir %x",page_dir);
    //对其size向上取整，得到需要分配的页数
    uint32_t page_size=up2(size,MEM_PAGE_SIZE)/MEM_PAGE_SIZE;
    vaddr=down2(vaddr,MEM_PAGE_SIZE);
    uint32_t current_addr=vaddr;

    for(int i=0;i<page_size;i++){
        uint32_t phy_addr=addr_alloc_page(&paddr_alloc,1);
        if(phy_addr<=0){
            //没有位置了
            log_printf("mem alloc no free page\n");
            return 0;
        }
        int err=memory_create_map((pde_t*)page_dir,current_addr,phy_addr,1,perm);
        //log_printf("create map from %x to  %x",current_addr,phy_addr);
        if(err<=0){
            // 创建映射失败
            log_printf("map memory fails\n");
            addr_free_page(&paddr_alloc,vaddr,i+1);
            return 0;
        }
        current_addr+=MEM_PAGE_SIZE;
        
    }
    //log_printf("finish alloc from %x vaddr to %x",vaddr,vaddr+size);
    return vaddr;

};


uint32_t memory_alloc_page (void){
    //分配一页内存, 给内核使用的，直接返回物理地址即可，前面作了0x0--0x8000 0000的映射
    uint32_t paddr=addr_alloc_page(&paddr_alloc, 1);
    return paddr;
};

static pde_t * get_task_current_page_dir(){
    //返回当前进程的一级页表地址

    task_t * current_task=get_task_current();
    return (pde_t *)current_task->tss.cr3;
}
void memory_free_page (uint32_t addr){
    //释放一页内存，内核和进程都可以用
    if(addr < MEMORY_TASK_BASE){
        //内核，直接释放
        addr_free_page(&paddr_alloc,addr,1);
    }else{
        //进程，需要找到对应的物理地址
        pte_t* current_pte=find_pte(get_task_current_page_dir(),addr,0);
        addr_free_page(&paddr_alloc,pte_paddr(current_pte),1);
        current_pte->v=0;
    }

};

void memory_destroy_uvm (uint32_t page_dir){
    //释放用户的页表
    
    uint32_t user_pde_index=pde_index(MEMORY_TASK_BASE);
    pde_t *user_pde=(pde_t*)page_dir+user_pde_index;

    //0x80000000之后的，就是进程自己的页表
    for(int i=user_pde_index;i<PDE_CNT;i++,user_pde++){
        if(!user_pde->present){
            continue;
        }

        pte_t *user_pte=(pte_t*)pde_paddr(user_pde);
        for(int j=0;j<PTE_CNT;j++){
            if(!user_pde->present){
                continue;
            }

            addr_free_page(&paddr_alloc, pte_paddr(user_pte), 1);
        }
        addr_free_page(&paddr_alloc, (uint32_t)pde_paddr(user_pde), 1);
    }
    // 页目录表
    addr_free_page(&paddr_alloc, page_dir, 1);
};
uint32_t memory_copy_uvm (uint32_t page_dir){
    // 复制一份page_dir 返回新页表的地址

    pde_t * dest_pde =(pde_t *)memory_create_uvm();
    if(!dest_pde){
        goto copy_uvm_failed;
    }



    // 这个是取操作系统最大内存的页表项, 之后的，就是需要复制的代码
    uint32_t user_pde_start=pde_index(MEMORY_TASK_BASE);
    //PDE_CNT是一级页表的最大项，4096/4
    pde_t *current_pde=(pde_t*)page_dir+user_pde_start;
    for(int i=user_pde_start;i<PDE_CNT;i++,current_pde++){
        if(current_pde->present==0){
            continue;//当前页表无效
        }
        pte_t *origin_pte=(pte_t *)pde_paddr(current_pde);
        for(int j=0;j<PTE_CNT;j++,origin_pte++){
            if (!origin_pte->present) {
                continue;
            }
                        // 分配物理内存
            uint32_t page = addr_alloc_page(&paddr_alloc, 1);
            //遍历二级页表的每一项
            pte_t *dest_pte=(pte_t *)addr_alloc_page(&paddr_alloc,1);
            uint32_t vaddr= (i << 22) | (j << 12); // 前10位，是pde在一级页表的索引，中间十位，是pte在二级页表的索引
            int err=memory_create_map(dest_pde,vaddr,page,1,get_pte_perm(origin_pte));
            if (err < 0) {
                goto copy_uvm_failed;
            }
            kernel_memcpy((void *)page,(void *)vaddr,MEM_PAGE_SIZE);
        }


    }
    return (uint32_t )dest_pde;
copy_uvm_failed:
    if (dest_pde) {
        memory_destroy_uvm((uint32_t)dest_pde);
    }
    return -1;
};

uint32_t memory_get_paddr (uint32_t page_dir, uint32_t vaddr) {
    //获取vaddr对应的物理地址，一级页表是page_dir
    pte_t* pte=find_pte((pde_t *)page_dir,vaddr,0);
    if(!pte){
        return 0;
    }
    //取vaddr的末12位，与 上页表的物理地址
    return pte_paddr(pte) | ( vaddr & (MEM_PAGE_SIZE-1) );
};



/**
 * @brief 在不同的进程空间中拷贝字符串
 * page_dir为目标页表，当前仍为老页表
 */
int memory_copy_uvm_data(uint32_t to, uint32_t page_dir, uint32_t from, uint32_t size) {
    char *buf, *pa0;

    while(size > 0){
        // 获取目标的物理地址, 也即其另一个虚拟地址
        uint32_t to_paddr = memory_get_paddr(page_dir, to);
        if (to_paddr == 0) {
            return -1;
        }

        // 计算当前可拷贝的大小
        uint32_t offset_in_page = to_paddr & (MEM_PAGE_SIZE - 1);
        uint32_t curr_size = MEM_PAGE_SIZE - offset_in_page;
        if (curr_size > size) {
            curr_size = size;       // 如果比较大，超过页边界，则只拷贝此页内的
        }

        kernel_memcpy((void *)to_paddr, (void *)from, curr_size);

        size -= curr_size;
        to += curr_size;
        from += curr_size;
  }

  return 0;
}


char * sys_sbrk(int incr){
    //在堆区申请incr块的内存大小, 返回原来堆起始的地址
    //只做堆区增加的情况
    ASSERT(incr >= 0); 

    //log_printf("call sbrk %x",incr);

    uint32_t unsigned_incr=(uint32_t) incr;

    task_t *current_task=get_task_current();
    uint32_t pre_heap_end=current_task->heap_end;

    // 如果地址为0，则返回有效的heap区域的顶端
    if (incr == 0) {
        log_printf("sbrk(0): end = 0x%x",current_task->heap_end );
        return (char *)current_task->heap_end ;
    } 
    uint32_t pre_incr=unsigned_incr;
    uint32_t end_offset=current_task->heap_end%MEM_PAGE_SIZE;

    if(end_offset){
        //当前堆还没有对齐
        uint32_t free_size=MEM_PAGE_SIZE - end_offset;
        if(free_size >=unsigned_incr){
            //申请的堆没到一页的大小
            current_task->heap_end+=unsigned_incr;
            unsigned_incr=0;
        }else{
            //先把一页end_offset的分配了
            unsigned_incr-=free_size;
            current_task->heap_end+=free_size;
        }
    }

    if(unsigned_incr){
        //剩下的交给内存管理，分配页表即可
        uint32_t err=memory_alloc_page_for(current_task->heap_end,unsigned_incr, PTE_P | PTE_U | PTE_W);
        if(err == 0){
            log_printf("sbrk: alloc mem failed.");
            return (char *)-1;
        }
        current_task->heap_end+=unsigned_incr;
    }
    //log_printf("sbrk(%d): end = 0x%x", pre_incr, current_task->heap_end);
    return (char * )pre_heap_end;
    

}


// /**
//  * @brief 调整堆的内存分配，返回堆之前的指针
//  */
// char * sys_sbrk(int incr) {
    
//     task_t * task = get_task_current();
//     char * pre_heap_end = (char * )task->heap_end;
//     int pre_incr = incr;

//     ASSERT(incr >= 0);

//     // 如果地址为0，则返回有效的heap区域的顶端
//     if (incr == 0) {
//         log_printf("sbrk(0): end = 0x%x", pre_heap_end);
//         return pre_heap_end;
//     } 
    
//     uint32_t start = task->heap_end;
//     uint32_t end = start + incr;

//     // 起始偏移非0
//     int start_offset = start % MEM_PAGE_SIZE;
//     if (start_offset) {
//         // 不超过1页，只调整
//         if (start_offset + incr <= MEM_PAGE_SIZE) {
//             task->heap_end = end;
//             return pre_heap_end;
//         } else {
//             // 超过1页，先只调本页的
//             uint32_t curr_size = MEM_PAGE_SIZE - start_offset;
//             start += curr_size;
//             incr -= curr_size;
//         }
//     }

//     // 处理其余的，起始对齐的页边界的
//     if (incr) {
//         uint32_t curr_size = end - start;
//         int err = memory_alloc_page_for(start, curr_size, PTE_P | PTE_U | PTE_W);
//         if (err < 0) {
//             log_printf("sbrk: alloc mem failed.");
//             return (char *)-1;
//         }
//     }

//     log_printf("sbrk(%d): end = 0x%x", pre_incr, end);
//     task->heap_end = end;
//     return (char * )pre_heap_end;        
// }
